// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/core/pretty_print.h>

#include <algorithm>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <iostream>
#include <limits>
#include <memory>
#include <sstream>  // IWYU pragma: keep
#include <string>
#include <string_view>
#include <type_traits>
#include <vector>

#include <nebula/core/array.h>
#include <nebula/core/chunked_array.h>
#include <nebula/core/record_batch.h>
#include <turbo/utility/status.h>
#include <nebula/core/table.h>
#include <nebula/types/type.h>
#include <nebula/types/type_traits.h>
#include <turbo/base/checked_cast.h>
#include <nebula/util/formatting.h>
#include <nebula/numeric/int_util_overflow.h>
#include <nebula/util/key_value_metadata.h>
#include <nebula/util/string.h>
#include <nebula/core/visit_array_inline.h>

namespace nebula {

    using internal::StringFormatter;
    using internal::ToChars;

    namespace {

        class PrettyPrinter {
        public:
            PrettyPrinter(const PrettyPrintOptions &options, std::ostream *sink)
                    : options_(options), indent_(options.indent), sink_(sink) {}

            inline void write(std::string_view data);

            inline void WriteIndented(std::string_view data);

            inline void Newline();

            inline void Indent();

            inline void IndentAfterNewline();

            void OpenArray(const Array &array);

            void CloseArray(const Array &array);

            void flush() { (*sink_) << std::flush; }

            PrettyPrintOptions ChildOptions(bool increment_indent = false) const {
                PrettyPrintOptions child_options = options_;
                if (increment_indent) {
                    child_options.indent = indent_ + child_options.indent_size;
                } else {
                    child_options.indent = indent_;
                }
                return child_options;
            }

        protected:
            const PrettyPrintOptions &options_;
            int indent_;
            std::ostream *sink_;
        };

        void PrettyPrinter::OpenArray(const Array &array) {
            if (!options_.skip_new_lines) {
                Indent();
            }
            (*sink_) << options_.array_delimiters.open;
            if (array.length() > 0) {
                Newline();
                indent_ += options_.indent_size;
            }
        }

        void PrettyPrinter::CloseArray(const Array &array) {
            if (array.length() > 0) {
                indent_ -= options_.indent_size;
                if (!options_.skip_new_lines) {
                    Indent();
                }
            }
            (*sink_) << options_.array_delimiters.close;
        }

        void PrettyPrinter::write(std::string_view data) { (*sink_) << data; }

        void PrettyPrinter::WriteIndented(std::string_view data) {
            Indent();
            write(data);
        }

        void PrettyPrinter::Newline() {
            if (options_.skip_new_lines) {
                return;
            }
            (*sink_) << "\n";
        }

        void PrettyPrinter::Indent() {
            for (int i = 0; i < indent_; ++i) {
                (*sink_) << " ";
            }
        }

        void PrettyPrinter::IndentAfterNewline() {
            if (options_.skip_new_lines) {
                return;
            }
            Indent();
        }

        class ArrayPrinter : public PrettyPrinter {
        public:
            ArrayPrinter(const PrettyPrintOptions &options, std::ostream *sink)
                    : PrettyPrinter(options, sink) {}

        private:
            template<typename FormatFunction>
            turbo::Status WriteValues(const Array &array, FormatFunction &&func,
                                      bool indent_non_null_values = true, bool is_container = false) {
                // `indent_non_null_values` should be false if `FormatFunction` applies
                // indentation itself.
                int window = is_container ? options_.container_window : options_.window;
                for (int64_t i = 0; i < array.length(); ++i) {
                    const bool is_last = (i == array.length() - 1);
                    // check if `length == 2 * window + 1` to eliminate ellipsis for only one element
                    if ((array.length() != 2 * window + 1) && (i >= window) &&
                        (i < (array.length() - window))) {
                        IndentAfterNewline();
                        (*sink_) << "...";
                        if (!is_last && options_.skip_new_lines) {
                            (*sink_) << options_.array_delimiters.element;
                        }
                        i = array.length() - window - 1;
                    } else if (array.is_null(i)) {
                        IndentAfterNewline();
                        (*sink_) << options_.null_rep;
                        if (!is_last) {
                            (*sink_) << options_.array_delimiters.element;
                        }
                    } else {
                        if (indent_non_null_values) {
                            IndentAfterNewline();
                        }
                        TURBO_RETURN_NOT_OK(func(i));
                        if (!is_last) {
                            (*sink_) << options_.array_delimiters.element;
                        }
                    }
                    Newline();
                }
                return turbo::OkStatus();
            }

            template<typename ArrayType, typename Formatter>
            turbo::Status WritePrimitiveValues(const ArrayType &array, Formatter *formatter) {
                auto appender = [&](std::string_view v) { (*sink_) << v; };
                auto format_func = [&](int64_t i) {
                    (*formatter)(array.get_view(i), appender);
                    return turbo::OkStatus();
                };
                return WriteValues(array, std::move(format_func));
            }

            template<typename ArrayType, typename T = typename ArrayType::TypeClass>
            turbo::Status WritePrimitiveValues(const ArrayType &array) {
                StringFormatter<T> formatter{array.type().get()};
                return WritePrimitiveValues(array, &formatter);
            }

            turbo::Status WriteValidityBitmap(const Array &array);

            turbo::Status PrintChildren(const std::vector<const Array *> &fields, int64_t offset,
                                        int64_t length) {
                for (size_t i = 0; i < fields.size(); ++i) {
                    write("\n");  // Always want newline before child array description
                    Indent();
                    std::stringstream ss;
                    ss << "-- child " << i << " type: " << fields[i]->type()->to_string() << "\n";
                    write(ss.str());

                    // Indent();
                    const Array *field = fields[i];
                    if (offset != 0) {
                        TURBO_RETURN_NOT_OK(
                                pretty_print(*field->slice(offset, length), ChildOptions(true), sink_));
                    } else {
                        TURBO_RETURN_NOT_OK(pretty_print(*field, ChildOptions(true), sink_));
                    }
                }
                return turbo::OkStatus();
            }

            //
            // WriteDataValues(): generic function to write values from an array
            //

            template<typename ArrayType, typename T = typename ArrayType::TypeClass>
            enable_if_has_c_type<T, turbo::Status> WriteDataValues(const ArrayType &array) {
                return WritePrimitiveValues(array);
            }

            turbo::Status WriteDataValues(const Fp16Array &array) {
                // XXX do not know how to format half floats yet
                StringFormatter<Int16Type> formatter{array.type().get()};
                return WritePrimitiveValues(array, &formatter);
            }

            template<typename ArrayType, typename T = typename ArrayType::TypeClass>
            enable_if_has_string_view<T, turbo::Status> WriteDataValues(const ArrayType &array) {
                return WriteValues(array, [&](int64_t i) {
                    if constexpr (T::is_utf8) {
                        (*sink_) << "\"" << array.get_view(i) << "\"";
                    } else {
                        (*sink_) << HexEncode(array.get_view(i));
                    }
                    return turbo::OkStatus();
                });
            }

            template<typename ArrayType, typename T = typename ArrayType::TypeClass>
            enable_if_decimal<T, turbo::Status> WriteDataValues(const ArrayType &array) {
                return WriteValues(array, [&](int64_t i) {
                    (*sink_) << array.format_value(i);
                    return turbo::OkStatus();
                });
            }

            template<typename ArrayType, typename T = typename ArrayType::TypeClass>
            enable_if_t<is_list_like_type<T>::value || is_list_view_type<T>::value, turbo::Status>
            WriteDataValues(const ArrayType &array) {
                const auto values = array.values();
                const auto child_options = ChildOptions();
                ArrayPrinter values_printer(child_options, sink_);

                return WriteValues(
                        array,
                        [&](int64_t i) {
                            // XXX this could be much faster if ArrayPrinter allowed specifying start and
                            // stop endpoints.
                            return values_printer.Print(
                                    *values->slice(array.value_offset(i), array.value_length(i)));
                        },
                        /*indent_non_null_values=*/false,
                        /*is_container=*/true);
            }

            turbo::Status WriteDataValues(const MapArray &array) {
                const auto keys = array.keys();
                const auto items = array.items();
                const auto child_options = ChildOptions();
                ArrayPrinter values_printer(child_options, sink_);

                return WriteValues(
                        array,
                        [&](int64_t i) {
                            IndentAfterNewline();
                            (*sink_) << "keys:";
                            Newline();
                            TURBO_RETURN_NOT_OK(values_printer.Print(
                                    *keys->slice(array.value_offset(i), array.value_length(i))));
                            Newline();
                            IndentAfterNewline();
                            (*sink_) << "values:";
                            Newline();
                            TURBO_RETURN_NOT_OK(values_printer.Print(
                                    *items->slice(array.value_offset(i), array.value_length(i))));
                            return turbo::OkStatus();
                        },
                        /*indent_non_null_values=*/false);
            }

        public:
            template<typename T>
            enable_if_t<std::is_base_of<PrimitiveArray, T>::value ||
                        std::is_base_of<FixedSizeBinaryArray, T>::value ||
                        std::is_base_of<BinaryArray, T>::value ||
                        std::is_base_of<LargeBinaryArray, T>::value ||
                        std::is_base_of<BinaryViewArray, T>::value ||
                        std::is_base_of<ListArray, T>::value ||
                        std::is_base_of<LargeListArray, T>::value ||
                        std::is_base_of<ListViewArray, T>::value ||
                        std::is_base_of<LargeListViewArray, T>::value ||
                        std::is_base_of<MapArray, T>::value ||
                        std::is_base_of<FixedSizeListArray, T>::value,
                    turbo::Status>
            Visit(const T &array) {
                turbo::Status st = array.validate();
                if (!st.ok()) {
                    (*sink_) << "<Invalid array: " << st.message() << ">";
                    return turbo::OkStatus();
                }

                OpenArray(array);
                if (array.length() > 0) {
                    TURBO_RETURN_NOT_OK(WriteDataValues(array));
                }
                CloseArray(array);
                return turbo::OkStatus();
            }

            turbo::Status Visit(const NullArray &array) {
                (*sink_) << array.length() << " nulls";
                return turbo::OkStatus();
            }

            turbo::Status Visit(const ExtensionArray &array) { return Print(*array.storage()); }

            turbo::Status Visit(const StructArray &array) {
                TURBO_RETURN_NOT_OK(WriteValidityBitmap(array));
                std::vector<const Array *> children;
                children.reserve(array.num_fields());
                for (int i = 0; i < array.num_fields(); ++i) {
                    children.emplace_back(array.field(i).get());
                }
                return PrintChildren(children, 0, array.length());
            }

            turbo::Status Visit(const UnionArray &array) {
                TURBO_RETURN_NOT_OK(WriteValidityBitmap(array));

                Newline();
                Indent();
                write("-- type_ids: ");
                UInt8Array type_codes(array.length(), array.type_codes(), nullptr, 0, array.offset());
                TURBO_RETURN_NOT_OK(pretty_print(type_codes, ChildOptions(true), sink_));

                if (array.mode() == UnionMode::DENSE) {
                    Newline();
                    Indent();
                    write("-- value_offsets: ");
                    Int32Array value_offsets(
                            array.length(), turbo::checked_cast<const DenseUnionArray &>(array).value_offsets(),
                            nullptr, 0, array.offset());
                    TURBO_RETURN_NOT_OK(pretty_print(value_offsets, ChildOptions(true), sink_));
                }

                // Print the children without any offset, because the type ids are absolute
                std::vector<const Array *> children;
                children.reserve(array.num_fields());
                for (int i = 0; i < array.num_fields(); ++i) {
                    children.emplace_back(array.field(i).get());
                }
                return PrintChildren(children, 0, array.length() + array.offset());
            }

            turbo::Status Visit(const DictionaryArray &array) {
                Newline();
                Indent();
                write("-- dictionary:\n");
                TURBO_RETURN_NOT_OK(pretty_print(*array.dictionary(), ChildOptions(true), sink_));

                Newline();
                Indent();
                write("-- indices:\n");
                return pretty_print(*array.indices(), ChildOptions(true), sink_);
            }

            turbo::Status Visit(const RunEndEncodedArray &array) {
                Newline();
                Indent();
                write("-- run_ends:\n");
                TURBO_RETURN_NOT_OK(pretty_print(*array.run_ends(), ChildOptions(true), sink_));

                Newline();
                Indent();
                write("-- values:\n");
                return pretty_print(*array.values(), ChildOptions(true), sink_);
            }

            turbo::Status Print(const Array &array) {
                if (array.device_type() != DeviceAllocationType::kCPU) {
                    // GH-43055: ideally we only copy start/end slices from non-CPU memory
                    // based on the window size that is being printed
                    TURBO_MOVE_OR_RAISE(auto array_cpu, array.copy_to(default_cpu_memory_manager()));
                    TURBO_RETURN_NOT_OK(visit_array_inline(*array_cpu, this));
                } else {
                    TURBO_RETURN_NOT_OK(visit_array_inline(array, this));
                }
                flush();
                return turbo::OkStatus();
            }
        };

        turbo::Status ArrayPrinter::WriteValidityBitmap(const Array &array) {
            Indent();
            write("-- is_valid:");

            if (array.null_count() > 0) {
                Newline();
                Indent();
                BooleanArray is_valid(array.length(), array.null_bitmap(), nullptr, 0,
                                      array.offset());
                return pretty_print(is_valid, ChildOptions(true), sink_);
            } else {
                write(" all not null");
                return turbo::OkStatus();
            }
        }

    }  // namespace

    turbo::Status pretty_print(const Array &arr, int indent, std::ostream *sink) {
        PrettyPrintOptions options;
        options.indent = indent;
        ArrayPrinter printer(options, sink);
        return printer.Print(arr);
    }

    turbo::Status pretty_print(const Array &arr, const PrettyPrintOptions &options,
                              std::ostream *sink) {
        ArrayPrinter printer(options, sink);
        return printer.Print(arr);
    }

    turbo::Status pretty_print(const Array &arr, const PrettyPrintOptions &options,
                              std::string *result) {
        std::ostringstream sink;
        TURBO_RETURN_NOT_OK(pretty_print(arr, options, &sink));
        *result = sink.str();
        return turbo::OkStatus();
    }

    turbo::Status pretty_print(const ChunkedArray &chunked_arr, const PrettyPrintOptions &options,
                              std::ostream *sink) {
        int num_chunks = chunked_arr.num_chunks();
        int indent = options.indent;
        int window = options.container_window;
        // Struct fields are always on new line
        bool skip_new_lines =
                options.skip_new_lines && (chunked_arr.type()->id() != Type::STRUCT);

        for (int i = 0; i < indent; ++i) {
            (*sink) << " ";
        }
        (*sink) << options.chunked_array_delimiters.open;
        if (!skip_new_lines) {
            *sink << "\n";
        }
        bool skip_element_delimiter = true;
        for (int i = 0; i < num_chunks; ++i) {
            if (skip_element_delimiter) {
                skip_element_delimiter = false;
            } else {
                (*sink) << options.chunked_array_delimiters.element;
                if (!skip_new_lines) {
                    *sink << "\n";
                }
            }
            if ((i >= window) && (i < (num_chunks - window))) {
                for (int i = 0; i < indent; ++i) {
                    (*sink) << " ";
                }
                (*sink) << "...";
                (*sink) << options.chunked_array_delimiters.element;
                if (!skip_new_lines) {
                    *sink << "\n";
                }
                i = num_chunks - window - 1;
                skip_element_delimiter = true;
            } else {
                PrettyPrintOptions chunk_options = options;
                chunk_options.indent += options.indent_size;
                ArrayPrinter printer(chunk_options, sink);
                TURBO_RETURN_NOT_OK(printer.Print(*chunked_arr.chunk(i)));
            }
        }
        if (!options.skip_new_lines) {
            *sink << "\n";
        }

        for (int i = 0; i < indent; ++i) {
            (*sink) << " ";
        }
        (*sink) << options.chunked_array_delimiters.close;

        return turbo::OkStatus();
    }

    turbo::Status pretty_print(const ChunkedArray &chunked_arr, const PrettyPrintOptions &options,
                              std::string *result) {
        std::ostringstream sink;
        TURBO_RETURN_NOT_OK(pretty_print(chunked_arr, options, &sink));
        *result = sink.str();
        return turbo::OkStatus();
    }

    turbo::Status pretty_print(const RecordBatch &batch, int indent, std::ostream *sink) {
        for (int i = 0; i < batch.num_columns(); ++i) {
            const std::string &name = batch.column_name(i);
            (*sink) << name << ": ";
            TURBO_RETURN_NOT_OK(pretty_print(*batch.column(i), indent + 2, sink));
            (*sink) << "\n";
        }
        (*sink) << std::flush;
        return turbo::OkStatus();
    }

    turbo::Status pretty_print(const RecordBatch &batch, const PrettyPrintOptions &options,
                              std::ostream *sink) {
        for (int i = 0; i < batch.num_columns(); ++i) {
            const std::string &name = batch.column_name(i);
            PrettyPrintOptions column_options = options;
            column_options.indent += 2;

            (*sink) << name << ": ";
            TURBO_RETURN_NOT_OK(pretty_print(*batch.column(i), column_options, sink));
            (*sink) << "\n";
        }
        (*sink) << std::flush;
        return turbo::OkStatus();
    }

    turbo::Status pretty_print(const Table &table, const PrettyPrintOptions &options,
                              std::ostream *sink) {
        TURBO_RETURN_NOT_OK(pretty_print(*table.schema(), options, sink));
        (*sink) << "\n";
        (*sink) << "----\n";

        PrettyPrintOptions column_options = options;
        column_options.indent += 2;
        for (int i = 0; i < table.num_columns(); ++i) {
            for (int j = 0; j < options.indent; ++j) {
                (*sink) << " ";
            }
            (*sink) << table.schema()->field(i)->name() << ":\n";
            TURBO_RETURN_NOT_OK(pretty_print(*table.column(i), column_options, sink));
            (*sink) << "\n";
        }
        (*sink) << std::flush;
        return turbo::OkStatus();
    }

    turbo::Status debug_print(const Array &arr, int indent) {
        return pretty_print(arr, indent, &std::cerr);
    }

    namespace {

        class SchemaPrinter : public PrettyPrinter {
        public:
            SchemaPrinter(const Schema &schema, const PrettyPrintOptions &options,
                          std::ostream *sink)
                    : PrettyPrinter(options, sink), schema_(schema) {}

            turbo::Status PrintType(const DataType &type, bool nullable);

            turbo::Status PrintField(const Field &field);

            void PrintVerboseMetadata(const KeyValueMetadata &metadata) {
                for (int64_t i = 0; i < metadata.size(); ++i) {
                    Newline();
                    Indent();
                    write(metadata.key(i) + ": '" + metadata.value(i) + "'");
                }
            }

            void PrintTruncatedMetadata(const KeyValueMetadata &metadata) {
                for (int64_t i = 0; i < metadata.size(); ++i) {
                    Newline();
                    Indent();
                    size_t size = metadata.value(i).size();
                    size_t truncated_size = std::max<size_t>(10, 70 - metadata.key(i).size() - indent_);
                    if (size <= truncated_size) {
                        write(metadata.key(i) + ": '" + metadata.value(i) + "'");
                        continue;
                    }

                    write(metadata.key(i) + ": '" + metadata.value(i).substr(0, truncated_size) +
                          "' + " + ToChars(size - truncated_size));
                }
            }

            void PrintMetadata(const std::string &metadata_type, const KeyValueMetadata &metadata) {
                if (metadata.size() > 0) {
                    Newline();
                    Indent();
                    write(metadata_type);
                    if (options_.truncate_metadata) {
                        PrintTruncatedMetadata(metadata);
                    } else {
                        PrintVerboseMetadata(metadata);
                    }
                }
            }

            turbo::Status Print() {
                for (int i = 0; i < schema_.num_fields(); ++i) {
                    if (i > 0) {
                        Newline();
                        Indent();
                    } else {
                        Indent();
                    }
                    TURBO_RETURN_NOT_OK(PrintField(*schema_.field(i)));
                }

                if (options_.show_schema_metadata && schema_.metadata() != nullptr) {
                    PrintMetadata("-- schema metadata --", *schema_.metadata());
                }
                flush();
                return turbo::OkStatus();
            }

        private:
            const Schema &schema_;
        };

        turbo::Status SchemaPrinter::PrintType(const DataType &type, bool nullable) {
            write(type.to_string());
            if (!nullable) {
                write(" not null");
            }
            for (int i = 0; i < type.num_fields(); ++i) {
                Newline();
                Indent();

                std::stringstream ss;
                ss << "child " << i << ", ";

                indent_ += options_.indent_size;
                WriteIndented(ss.str());
                TURBO_RETURN_NOT_OK(PrintField(*type.field(i)));
                indent_ -= options_.indent_size;
            }
            return turbo::OkStatus();
        }

        turbo::Status SchemaPrinter::PrintField(const Field &field) {
            write(field.name());
            write(": ");
            TURBO_RETURN_NOT_OK(PrintType(*field.type(), field.nullable()));

            if (options_.show_field_metadata && field.metadata() != nullptr) {
                indent_ += options_.indent_size;
                PrintMetadata("-- field metadata --", *field.metadata());
                indent_ -= options_.indent_size;
            }
            return turbo::OkStatus();
        }

    }  // namespace

    turbo::Status pretty_print(const Schema &schema, const PrettyPrintOptions &options,
                              std::ostream *sink) {
        SchemaPrinter printer(schema, options, sink);
        return printer.Print();
    }

    turbo::Status pretty_print(const Schema &schema, const PrettyPrintOptions &options,
                              std::string *result) {
        std::ostringstream sink;
        TURBO_RETURN_NOT_OK(pretty_print(schema, options, &sink));
        *result = sink.str();
        return turbo::OkStatus();
    }

}  // namespace nebula
